#!/usr/bin/env python3
import subprocess
import sys
import os
import re
from collections import defaultdict

class OSCAPRuleChecker:
    def __init__(self, profile="basic", ruleset="tencentos4"):
        self.profile = profile
        self.ruleset = ruleset
        self.ds_file = f"/usr/share/xml/scap/ssg/content/ssg-{ruleset}-ds.xml"
        self.profile_id = f"xccdf_org.ssgproject.content_profile_{profile}"
        
        # 统计结果
        self.results = {
            'PASS': 0,
            'FAIL': 0,
            'NOTAPPLICABLE': 0,
            'ERROR': 0
        }
        
        self.failed_rules = []
        self.missing_rules = []
        self.remediated_rules = []

    def run_oscap_rule(self, rule_name):
        """执行单个规则的oscap检查"""
        rule_id = f"xccdf_org.ssgproject.content_rule_{rule_name}"
        
        cmd = [
            "sudo", "oscap", "xccdf", "eval",
            "--profile", self.profile_id,
            "--rule", rule_id,
            self.ds_file
        ]
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            return self.parse_oscap_output(result.stdout, result.stderr, result.returncode)
        except subprocess.TimeoutExpired:
            print(f"  超时: {rule_name}")
            return "ERROR"
        except Exception as e:
            print(f"  执行错误: {rule_name} - {e}")
            return "ERROR"

    def parse_oscap_output(self, stdout, stderr, returncode):
        """解析oscap输出结果"""
        if "pass" in stdout.lower():
            return "PASS"
        elif "fail" in stdout.lower():
            return "FAIL"
        elif "notapplicable" in stdout.lower() or "not applicable" in stdout.lower():
            return "NOTAPPLICABLE"
        elif "notchecked" in stdout.lower() or "not checked" in stdout.lower():
            return "NOTAPPLICABLE"
        elif "error" in stderr.lower() or returncode != 0:
            return "ERROR"
        else:
            return "ERROR"

    def find_remediation_script(self, rule_name, section):
        """查找对应的修复脚本"""
        # 根据章节号构建可能的路径
        section_parts = section.split('.')
        if len(section_parts) >= 2:
            # 构建章节路径，如 1.1.1 -> 1-initial-setup/1.1-filesystem-configuration
            chapter = section_parts[0]
            subsection = f"{section_parts[0]}.{section_parts[1]}"
            
            # 尝试多种可能的路径格式
            possible_paths = [
                f"remediations/{chapter}-*/{subsection}-*/{section}-*.sh",
                f"remediations/{chapter}-*/{subsection}-*/*{rule_name}*.sh",
            ]
            
            # 使用glob模式匹配
            import glob
            for pattern in possible_paths:
                matches = glob.glob(pattern)
                if matches:
                    return matches[0]  # 返回第一个匹配的文件
        
        # 递归搜索包含规则名称的脚本
        for root, dirs, files in os.walk("remediations"):
            for file in files:
                if file.endswith('.sh'):
                    # 检查文件名是否包含规则名称或章节号
                    if rule_name in file or section in file:
                        return os.path.join(root, file)
                    
                    # 检查文件名是否以章节号开头
                    if file.startswith(section):
                        return os.path.join(root, file)
        
        return None

    def run_remediation(self, rule_name, section):
        """执行修复脚本"""
        script_path = self.find_remediation_script(rule_name, section)
        
        if not script_path:
            print(f"    未找到修复脚本: {rule_name}")
            return False
        
        try:
            print(f"    执行修复脚本: {script_path}")
            result = subprocess.run(["sudo", "bash", script_path], 
                                  capture_output=True, text=True, timeout=60)
            
            if result.returncode == 0:
                print(f"    修复成功: {rule_name}")
                return True
            else:
                print(f"    修复失败: {rule_name} - {result.stderr}")
                return False
                
        except Exception as e:
            print(f"    修复执行错误: {rule_name} - {e}")
            return False

    def check_rules_from_file(self, rules_file):
        """从文件读取规则并检查"""
        rules_dict = self.parse_rules_file(rules_file)
        
        print(f"开始检查 {len(sum(rules_dict.values(), []))} 个规则...")
        print("=" * 60)
        
        for section, rules in rules_dict.items():
            print(f"\n检查章节: {section}")
            print("-" * 40)
            
            for rule in rules:
                print(f"检查规则: {rule}")
                result = self.run_oscap_rule(rule)
                
                if result == "ERROR":
                    self.missing_rules.append(rule)
                    print(f"  结果: 规则未找到")
                    continue
                
                self.results[result] += 1
                print(f"  结果: {result}")
                
                # 如果失败，尝试修复
                if result == "FAIL":
                    self.failed_rules.append((rule, section))
                    print(f"  尝试修复...")
                    
                    if self.run_remediation(rule, section):
                        # 重新检查
                        print(f"  重新检查规则: {rule}")
                        recheck_result = self.run_oscap_rule(rule)
                        print(f"  修复后结果: {recheck_result}")
                        
                        if recheck_result in ["PASS", "NOTAPPLICABLE"]:
                            self.results["FAIL"] -= 1
                            self.results[recheck_result] += 1
                            self.remediated_rules.append(rule)
                            print(f"  修复成功!")
                        else:
                            print(f"  修复后仍然失败")

    def parse_rules_file(self, rules_file):
        """解析规则文件"""
        rules_dict = defaultdict(list)
        current_section = None
        
        with open(rules_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                
                # 检查是否是章节标题 (如 1.1.1)
                if re.match(r'^\d+\.\d+\.\d+', line):
                    current_section = line
                # 检查是否是规则 (以 - 开头)
                elif line.startswith('- ') and current_section:
                    rule = line[2:].strip()
                    rules_dict[current_section].append(rule)
        
        return rules_dict

    def print_summary(self):
        """打印检查结果汇总"""
        print("\n" + "=" * 60)
        print("检查结果汇总")
        print("=" * 60)
        
        total = sum(self.results.values())
        print(f"总规则数: {total}")
        print(f"通过 (PASS): {self.results['PASS']}")
        print(f"不适用 (NOTAPPLICABLE): {self.results['NOTAPPLICABLE']}")
        print(f"失败 (FAIL): {self.results['FAIL']}")
        print(f"错误/未找到: {self.results['ERROR']}")
        
        if self.remediated_rules:
            print(f"\n成功修复的规则 ({len(self.remediated_rules)}):")
            for rule in self.remediated_rules:
                print(f"  - {rule}")
        
        if self.missing_rules:
            print(f"\n未找到的规则 ({len(self.missing_rules)}):")
            for rule in self.missing_rules:
                print(f"  - {rule}")
        
        # 计算成功率
        success_count = self.results['PASS'] + self.results['NOTAPPLICABLE']
        if total > 0:
            success_rate = (success_count / total) * 100
            print(f"\n总体成功率: {success_rate:.1f}%")

def main():
    if len(sys.argv) < 2:
        print("用法: python3 oscap_rule_checker.py <rules_file> [profile] [ruleset]")
        print("示例: python3 oscap_rule_checker.py rules.txt basic tencentos4")
        sys.exit(1)
    
    rules_file = sys.argv[1]
    profile = sys.argv[2] if len(sys.argv) > 2 else "basic"
    ruleset = sys.argv[3] if len(sys.argv) > 3 else "tencentos4"
    
    if not os.path.exists(rules_file):
        print(f"错误: 规则文件 {rules_file} 不存在")
        sys.exit(1)
    
    checker = OSCAPRuleChecker(profile, ruleset)
    checker.check_rules_from_file(rules_file)
    checker.print_summary()

if __name__ == "__main__":
    main()
